﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;


#if X64
    using IMAGE_OPTIONAL_HEADER = GypeLaunch.Wrapper.IMAGE_OPTIONAL_HEADER64;
    using IMAGE_NT_HEADERS = GypeLaunch.Wrapper.IMAGE_NT_HEADERS64;
    using IMAGE_THUNK_DATA = GypeLaunch.Wrapper.IMAGE_THUNK_DATA64;
#else
using IMAGE_OPTIONAL_HEADER = GypeLaunch.Wrapper.IMAGE_OPTIONAL_HEADER32;
    using IMAGE_NT_HEADERS = GypeLaunch.Wrapper.IMAGE_NT_HEADERS32;
    using IMAGE_THUNK_DATA = GypeLaunch.Wrapper.IMAGE_THUNK_DATA32;
#endif

namespace GypeLaunch.Wrapper
{
    public class Module
    {
        public Win32.ModuleEntry32 ModuleEntry { get; set; }
        public Process Process { get; set; }
        public IMAGE_DOS_HEADER DosHeader { get; private set; }
        public IMAGE_NT_HEADERS NtHeader { get; set; }

        public Module()
        {
        }

        public Module(Win32.ModuleEntry32 module)
        {
            ModuleEntry = module;
        }

        public Module(Process process, IntPtr module)
        {
            Win32.EnumProcessModules(process.Id, delegate(Win32.ModuleEntry32 mod)
            {
                if (mod.Module == module)
                {
                    ModuleEntry = mod;
                    return true;
                }
                return false;
            });
        }

        public static List<Module> GetModules(Process process)
        {
            List<Module> mods = new List<Module>();

            foreach (Win32.ModuleEntry32 module in Win32.EnumProcessModules(process.Id))
                mods.Add(new Module()
                    {
                        ModuleEntry = module,
                        Process = process
                    });

            return mods;
        }

        public static Module FromLibrary(string path)
        {
            var handle = Win32.LoadLibrary(path);
            return new Module(Process.GetCurrentProcess(), handle);
        }

        public void ReadPortableExecutableHeader()
        {
            var handle = Win32.OpenProcess(Win32.ProcessAccessRights.ALL_ACCESS, false, Process.Id);

            DosHeader = Win32.Read<IMAGE_DOS_HEADER>(handle, ModuleEntry.BaseAddress);
            var ntAddress = ModuleEntry.BaseAddress.ToInt64() + DosHeader.lfanew;
            NtHeader = Win32.Read<IMAGE_NT_HEADERS>(handle, new IntPtr(ntAddress));

            var imageSectionHeaderAddress = ntAddress + 24 + NtHeader.FileHeader.SizeOfOptionalHeader;
            var imageSectionHeader = Win32.Read<IMAGE_SECTION_HEADER>(handle, new IntPtr(imageSectionHeaderAddress));

            var name = Encoding.ASCII.GetString(imageSectionHeader.Name);



            //IAT Code
            var IMAGE_DIRECTORY_ENTRY_IMPORT = 0;
            var dataDir = NtHeader.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];

            var importTableAddress = dataDir.VirtualAddress + ModuleEntry.BaseAddress.ToInt64();
            var importTable = Win32.Read<IMAGE_IMPORT_DESCRIPTOR>(handle, new IntPtr(importTableAddress));

            for (; importTable.Characteristics != 0; importTableAddress += Marshal.SizeOf(importTable))
            {
                importTable = Win32.Read<IMAGE_IMPORT_DESCRIPTOR>(handle, new IntPtr(importTableAddress));

                if (importTable.FirstThunk == 0)
                    continue;

                var nname = Win32.ReadString(handle, new IntPtr(importTable.Name));

                var thunk = Win32.Read<IMAGE_THUNK_DATA>(handle, new IntPtr(importTable.FirstThunk + ModuleEntry.BaseAddress.ToInt64()));
                var origThunk = Win32.Read<IMAGE_THUNK_DATA>(handle, new IntPtr(importTable.Characteristics + ModuleEntry.BaseAddress.ToInt64()));

                for (int i = 0; origThunk.u1 != 0; ++i)
                {
                    var address = importTable.FirstThunk + ModuleEntry.BaseAddress.ToInt64() + i * Marshal.SizeOf(thunk);
                    Win32.Read<IMAGE_THUNK_DATA>(handle, new IntPtr(address));

                    address = importTable.Characteristics + ModuleEntry.BaseAddress.ToInt64() + i * Marshal.SizeOf(thunk);
                    Win32.Read<IMAGE_THUNK_DATA>(handle, new IntPtr(address));

                    var importAddress = new IntPtr(origThunk.u1 + ModuleEntry.BaseAddress.ToInt64());
                    IMAGE_IMPORT_BY_NAME import = new IMAGE_IMPORT_BY_NAME();
                    import.Hint = Win32.Read<UInt16>(handle, importAddress);

                    var importName = Win32.ReadString(handle, importAddress);
                }
            }
            
            Win32.CloseHandle(handle);
        }
    }
}
